#include "mpi_partitionings.h"

void partition_static_cycles_func (float *samples, ConfigData *data) {

  // Variables from the data structure
  unsigned int height       = data->height;
  unsigned int width        = data->width;
  unsigned int strip_height = data->strip_height;      
  unsigned int num_procs    = data->mpi_procs;
  unsigned int rank         = data->mpi_rank;
  
  // Variables used in the computations
  unsigned int start_width  = 0;
  unsigned int end_width    = width - 1;
  unsigned int start_height = 0;
  unsigned int x, y, i;

  // First node gets first strip_height rows, etc.
  for (x = start_width; x <= end_width; x++) {
    start_height = rank * strip_height;
    while (start_height < height) {
      for (y = start_height; y < start_height + strip_height; y++) {
        if (y < height) {
          i = (y * width + x) * 3;
          shadePixel( samples + i, x, y, data );
        }
      }
      start_height = start_height + num_procs*strip_height;
    }
  }
}

void partition_static_strips_func (float *samples, ConfigData *data) {

  // Variables from the data structure
  unsigned int height       = data->height;
  unsigned int width        = data->width;
  unsigned int num_procs    = data->mpi_procs;
  unsigned int rank         = data->mpi_rank;
  
  // Variables used in the computations
  unsigned int x, y, i;
  
  // Take into account non-nicely-divisable height / num_procs
  unsigned int end_height   = (rank == (num_procs-1)) ? height - 1 : (rank+1) * ceil(height/num_procs) - 1;
  unsigned int start_height = rank * ceil(height/num_procs);
                                

  // Loop through start_width -> end_width, start_height -> end_height
  for (x = 0; x < width; x++) {
    for (y = start_height; y <= end_height; y++) {
      i = (y * width + x) * 3;
      shadePixel( samples + i, x, y, data );
    }
  }
}

int get_size (unsigned int rank, ConfigData *data, task_list **my_tasks) {
  // Variables from the data structure
  unsigned int height       = data->height;
  unsigned int width        = data->width;
  unsigned int num_procs    = data->mpi_procs;
  
  unsigned int even_num_procs = num_procs - (num_procs % 2);
  
  // Variables used in the computations
  unsigned int pow_2;
  
  unsigned int start_height, end_height;
  unsigned int start_width,  end_width;
  
  unsigned int horiz_splits, vert_splits;
  
  unsigned int height_of_strip = 0;
  
  // Find the largest power of 2 inside of (#proc (-1 if odd))...
  // That is the number of veritcal splits. The divisor is the number of 
  // horizontal splits.

  // Get vert splits
  vert_splits = 0; 
  for (pow_2 = 2; pow_2 <= even_num_procs; pow_2 = pow_2 * 2) {
    if (even_num_procs % pow_2 == 0) {
      vert_splits = pow_2;
    }
  }

  // Get horiz splits
  horiz_splits = even_num_procs / vert_splits;
  
  // Get Height of bottom strip, if num_procs is odd
  if (num_procs % 2) {
    height_of_strip = round((height / (float)(vert_splits*horiz_splits)) /
                        (1 + (1 / (float)(vert_splits*horiz_splits))));
  } else {
    height_of_strip = 0;
  }
    
  // Get the (start, end) (width, height)
  if ((num_procs % 2) && (rank == (num_procs - 1))) {
    // If num_procs is odd, then make the last processor into a long strip  
    start_width  = 0;
    end_width    = width - 1;
    end_height   = height - 1;
    start_height = end_height - height_of_strip - 1;
    
  } else {
    start_width  = (rank % vert_splits) * (width / vert_splits);
    end_width    = start_width + (width / vert_splits) - 1;
    
    start_height = (rank / vert_splits) * (height - height_of_strip) / horiz_splits;
    end_height   = start_height + (height - height_of_strip) / horiz_splits - 1;
  }

  // Save to the memory location at **my_tasks
  *my_tasks = create_task_list(start_height, start_width, end_height, end_width);
  
  printf ("(%d) H%d-%d, W%d-%d [%f]\n", rank, start_height, end_height, start_width, end_width, (height / (float)(vert_splits*horiz_splits)) / (1 + (1 / (float)(vert_splits*horiz_splits))));
    
  return (end_height-start_height+1) * (end_width-start_width+1) * 3;
  
}

void partition_static_blocks_func (float *samples, ConfigData *data, task_list **my_tasks) {
  unsigned int node  = data->mpi_rank;
  int * recv_arr_int = (int*) malloc (4*sizeof(int));
  
  get_size (node, data, my_tasks);
  get_array(*my_tasks, recv_arr_int);
  
  unsigned int start_height = recv_arr_int[0];
  unsigned int start_width  = recv_arr_int[1];
  unsigned int end_height   = recv_arr_int[2];
  unsigned int end_width    = recv_arr_int[3];

  unsigned int x, y, i;
  
  unsigned int width  = data->width;
  
  free (recv_arr_int);
  
  //printf ("(%d) Starting computations (H%d-%d, W%d-%d)\n", node, start_height, end_height, start_width, end_width);
  
  // Loop through start_width -> end_width, start_height -> end_height
  for (x = start_width; x <= end_width; x++) {
    for (y = start_height; y <= end_height; y++) {
      i = (y * width + x) * 3;
      shadePixel( samples + i, x, y, data );
    }
  }

}
